import torch
import joblib
import numpy as np
import json
from models.price_predictor import PricePredictor

NUMERIC_COLUMNS = ["预算", "流行度", "平均评分", "评分人数", "已上映天数", "入座率", "影院评分"]


def predict_price_from_features(feature_dict, model_path="checkpoints/best_model.pt",
                                config_path="config/model_config.json",
                                scaler_path="data/processed/scaler.pkl"):
    for key in NUMERIC_COLUMNS:
        if key not in feature_dict:
            raise ValueError(f"缺少特征: {key}")

    # 构造输入
    input_vals = np.array([feature_dict[k] for k in NUMERIC_COLUMNS]).reshape(1, -1)
    scaler = joblib.load(scaler_path)
    input_scaled = scaler.transform(input_vals)
    x = torch.tensor(input_scaled, dtype=torch.float32)

    # 加载模型结构
    with open(config_path, "r") as f:
        model_cfg = json.load(f)["model"]

    input_dim = x.shape[1]
    model = PricePredictor(config_path, input_dim)
    model.load_state_dict(torch.load(model_path, weights_only=True))
    model.eval()

    with torch.no_grad():
        output = model(x)
        return output.item()


if __name__ == "__main__":
    example = {
        "预算": 237000000,
        "流行度": 169.4,
        "平均评分": 8.4,
        "评分人数": 196839,
        "已上映天数": 23,
        "入座率": 70.5,
        "影院评分": 4
    }
    result = predict_price_from_features(example)
    print(f"🎯 预测票价: {result:.2f} 元")
